import numpy as np
import faiss

class FaissKNeighbors:
    def __init__(self, k=1, res=None):
        self.index = None
        self.y = None
        self.k = k
        self.res = res

    def fit(self, X, y):
        # 初始化L2距离索引
        d = X.shape[1]  # 向量维度
        if self.res is not None:
            # 使用GPU索引
            self.index = faiss.GpuIndexFlatL2(self.res, d)
        else:
            # 使用CPU索引
            self.index = faiss.IndexFlatL2(d)
        
        self.index.add(X.astype(np.float32))
        self.y = y.copy()

    def predict(self, X):
        distances, indices = self.index.search(X.astype(np.float32), self.k)
        votes = self.y[indices]
        predictions = np.array([np.argmax(np.bincount(vote)) for vote in votes])
        return predictions

    def score(self, X, y):
        predictions = self.predict(X)
        return np.mean(predictions == y)
